package weiyao.Detector;

import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import us.codecraft.webmagic.selector.Html;
import weiyao.GUI.ScanResultCallback;
import weiyao.SiteManager;
import weiyao.Vul.VulnerabilityDetector;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public abstract class BaseDetector implements VulnerabilityDetector {
    protected SiteManager siteManager;

    public void setSiteManager(SiteManager siteManager) {
        this.siteManager = siteManager;
    }

    protected void testFormParameters(String url, String method, Map<String, String> params,
                                      String[] payloads, VulnerabilityTester tester,
                                      ScanResultCallback callback) {

        testFormParameters(url, method, params, payloads, tester, callback, "");
    }
    protected Map<String, String> collectFormParams(Html form) {
        Map<String, String> formParams = new HashMap<>();
        List<String> inputElements = form.xpath("//input").all();

        for (String inputHtml : inputElements) {
            Html input = new Html(inputHtml);
            String name = input.xpath("//input/@name").get();
            String value = input.xpath("//input/@value").get();

            if (name != null && !name.isEmpty()) {
                formParams.put(name, value != null ? value : "");
            }
        }
        return formParams;
    }

    protected void testFormParameters(String url, String method, Map<String, String> params,
                                      String[] payloads, VulnerabilityTester tester,
                                      ScanResultCallback callback, String vulnType) {
        for (String paramName : params.keySet()) {
            for (String payload : payloads) {
                Map<String, String> testParams = new HashMap<>(params);
                testParams.put(paramName, payload);

                try {
                    String response = sendHttpRequest(url, method, testParams);
                    if (tester.testVulnerability(url, paramName, payload, response)) {
                        // 生成格式化的漏洞详情
                        String details = generateVulnerabilityDetails(
                                vulnType, url, method, paramName, payload, response);

                        callback.onVulnerabilityFound(vulnType, details);
                    }
                } catch (Exception e) {
                    callback.onError("请求失败: " + e.getMessage());
                }
            }
        }
    }

    // 生成格式化的漏洞详情字符串
    protected String generateVulnerabilityDetails(String type, String url, String method,
                                                  String paramName, String payload, String response) {
        StringBuilder details = new StringBuilder();
        details.append("漏洞类型: ").append(type).append("\n");
        details.append("URL: ").append(url).append("\n");
        details.append("方法: ").append(method).append("\n");

        if (paramName != null) {
            details.append("参数: ").append(paramName).append("\n");
        }

        if (payload != null) {
            details.append("载荷: ").append(payload).append("\n");
        }

        // 添加响应中的关键信息
        if (response != null && !response.isEmpty()) {
            details.append("响应摘要: ").append(extractInjectionResult(response));

        }

        return details.toString();
    }
    public static String extractInjectionResult(String htmlResponse) {
        // 定义匹配模式
        Pattern pattern = Pattern.compile(
                "<font size='5' color= '#99FF00'>(.*?)</font>",
                Pattern.DOTALL
        );

        Matcher matcher = pattern.matcher(htmlResponse);

        if (matcher.find()) {
            String resultContent = matcher.group(1);

            // 清理HTML标签并提取核心信息
            return resultContent
                    .replaceAll("<br>", "\n")  // 转换换行符
                    .replaceAll("</?[^>]+>", "")  // 移除所有HTML标签
                    .replaceAll("\\s+", " ")  // 压缩多余空格
                    .trim();
        }
        return "No injection result found";
    }
    protected String sendHttpRequest(String url, String method, Map<String, String> params) throws Exception {
        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
            HttpUriRequest request;

            if ("POST".equalsIgnoreCase(method)) {
                // 构建 POST 请求
                HttpPost post = new HttpPost(url);
                StringBuilder paramStr = new StringBuilder();
                for (Map.Entry<String, String> entry : params.entrySet()) {
                    if (paramStr.length() > 0) paramStr.append("&");
                    // 对参数值进行 URL 编码
                    String encodedValue = URLEncoder.encode(entry.getValue(), StandardCharsets.UTF_8.toString());
                    paramStr.append(entry.getKey()).append("=").append(encodedValue);
                }
                post.setEntity(new StringEntity(paramStr.toString()));
                post.setHeader("Content-Type", "application/x-www-form-urlencoded");
                request = post;
            } else {
                // 构建 GET 请求
                StringBuilder urlWithParams = new StringBuilder(url);
                boolean hasQuery = url.contains("?");

                for (Map.Entry<String, String> entry : params.entrySet()) {
                    // 对参数值进行 URL 编码
                    String encodedValue = URLEncoder.encode(entry.getValue(), StandardCharsets.UTF_8.toString());
                    urlWithParams.append(hasQuery ? "&" : "?")
                            .append(entry.getKey())
                            .append("=")
                            .append(encodedValue);
                    hasQuery = true;
                }
                request = new HttpGet(urlWithParams.toString());
            }

            // 添加自定义请求头
            siteManager.getCurrentSite().getHeaders().forEach(request::addHeader);
            HttpResponse response = httpClient.execute(request);
            return EntityUtils.toString(response.getEntity());
        }
    }

    // 在 BaseDetector 中添加智能检测方法
    protected boolean isXssVulnerable(String response, String payload) {
        String decodedPayload = java.net.URLDecoder.decode(payload, StandardCharsets.UTF_8);

        // 1. 检查原始 payload 是否直接出现
        if (response.contains(payload) || response.contains(decodedPayload)) {
            return true;
        }

        // 2. 检查常见编码变体
        String htmlEncoded = payload.replace("<", "&lt;").replace(">", "&gt;");
        if (response.contains(htmlEncoded)) {
            // 如果服务器转义了特殊字符，说明有防护
            return false;
        }

        // 3. 检查关键片段 (更智能的检测)
        if (payload.contains("alert")) {
            // 检查是否出现 alert 关键字
            if (response.contains("alert")) {
                return true;
            }
            // 检查混淆后的 alert
            if (response.contains("al"+"ert") || response.contains("al/x65rt")) {
                return true;
            }
        }

        // 4. 检查 payload 的关键部分
        String[] keyFragments = {"script", "onerror", "onload", "svg", "iframe"};
        for (String fragment : keyFragments) {
            if (payload.contains(fragment) && response.contains(fragment)) {
                return true;
            }
        }

        return false;
    }

    protected String makeAbsoluteUrl(String baseUrl, String relativeUrl) {
        if (relativeUrl == null) return baseUrl;
        if (relativeUrl.startsWith("http")) return relativeUrl;

        if (relativeUrl.startsWith("/")) {
            int protocolEnd = baseUrl.indexOf("://") + 3;
            int domainEnd = baseUrl.indexOf("/", protocolEnd);
            if (domainEnd == -1) return baseUrl + relativeUrl;
            return baseUrl.substring(0, domainEnd) + relativeUrl;
        }

        int lastSlash = baseUrl.lastIndexOf("/");
        return baseUrl.substring(0, lastSlash + 1) + relativeUrl;
    }

    @FunctionalInterface
    protected interface VulnerabilityTester {
        boolean testVulnerability(String url, String paramName, String payload, String response);
    }




}